# -*- coding: utf-8 -*-
"""
    Project name：code_potentialflow
# -------------------------------
    File name：rect_sosi_flow.py
    Created on：2025/2/24 21:39
    Author：(input)
    Description:
    矩形区域内源汇流动
"""
import numpy as np
import scipy.special as sp
from scipy.optimize import root_scalar
from geometry import Grid
from visualize import SosiFlowPlotter, GridPlotter

from numpy.polynomial import Polynomial

from parasolver import ParameterSolver

# 源汇组合类（一源一汇）
class SourceSink:
    """
    定义室内一源一汇流动（Source-sink Flow）的基本计算。

    该类用于计算势流的复势函数、流函数、速度分布等参数。
    """

    def __init__(self, sosi_strength, source_org, sink_org, C, k, length, height):
        """
        初始化一源一汇流动参数

        Parameters
        ----------
        sosi_strength : tuple or list
            Strength of the source and sink pair (source_strength, sink_strength) in m^2/s
        source_org : tuple or complex
            Source origin coordinates (length_1, height_1) in meters
        sink_org : tuple or complex
            Sink origin coordinates (length_2, height_2) in meters
        C : float
            Parameter for the Schwarz-Christoffel (SC) transformation function
        k : float
            Modulus number for elliptic functions
        length : float
            Length of room (m)
        height : float
            Height of room (m)
        """
        # 参数验证
        if not isinstance(sosi_strength, (tuple, list)) or len(sosi_strength) != 2:
            raise ValueError("sosi_strength must be a tuple or list with 2 elements (source, sink)")
        if C is None or not isinstance(C, (int, float)):
            raise ValueError("C must be a valid number (int or float), not None")
        if k is None or not isinstance(k, (int, float)):
            raise ValueError("k must be a valid number (int or float), not None")
        if length <= 0:
            raise ValueError("房间长度 length 必须大于 0")
        if height <= 0:
            raise ValueError("房间高度 height 必须大于 0")

        # 源和汇的强度
        self.Q_source = sosi_strength[0]
        self.Q_sink = sosi_strength[1]

        # 房间尺寸
        self.L = length
        self.H = height

        # 源和汇的坐标
        if isinstance(source_org, (tuple, list)) and len(source_org) == 2:
            self.L1, self.H1 = source_org
        elif isinstance(source_org, complex):
            self.L1, self.H1 = source_org.real, source_org.imag
        else:
            raise ValueError("source_org must be a tuple/list (x, y) or complex number")
        if isinstance(sink_org, (tuple, list)) and len(sink_org) == 2:
            self.L2, self.H2 = sink_org
        elif isinstance(sink_org, complex):
            self.L2, self.H2 = sink_org.real, sink_org.imag
        else:
            raise ValueError("sink_org must be a tuple/list (x, y) or complex number")

        # 检查坐标是否在房间范围内
        if not (0 <= self.L1 <= self.L and 0 <= self.H1 <= self.H):
            raise ValueError("Source coordinates must be within room dimensions")
        if not (0 <= self.L2 <= self.L and 0 <= self.H2 <= self.H):
            raise ValueError("Sink coordinates must be within room dimensions")

        # SC 变换参数
        self.C0 = float(C)
        self.k0 = float(k)
        self.m0 = self.k0 ** 2
        self.k0p = np.sqrt(1 - self.m0)
        self.m0p = 1 - self.m0
        self.epsilon = 1e-10

        # 初始化并计算无量纲坐标
        self.z_so = None
        self.z_si = None
        self.compute_z_coordinates()
        if self.z_so is None or self.z_si is None:
            raise ValueError("Failed to compute z_so or z_si during initialization")

    def __repr__(self):
        """
        返回类的字符串表示，方便打印查看
        """
        return (f"SourceSink(source_strength={self.Q_source}, sink_strength={self.Q_sink}, "
                f"source_org=({self.L1}, {self.H1}), sink_org=({self.L2}, {self.H2}), "
                f"C={self.C0}, k={self.k0}, length={self.L}, height={self.H})")

    def compute_z_coordinates(self):
        '''
        计算源点和汇点的无量纲坐标，并保存到类变量 z_so 和 z_si 中
        '''
        E = self.L1 - self.L / 2 + 1j * self.H1  # 源点
        F = self.L2 - self.L / 2 + 1j * self.H2  # 汇点
        self.z_so = E / self.H  # 源点无量纲坐标
        self.z_si = F / self.H  # 汇点无量纲坐标
        if self.z_so is None or self.z_si is None:
            raise ValueError("z_so or z_si computation resulted in None")

    def _SCT_g(self, X, Y):
        """
        计算复数域的椭圆正弦函数 sn(z, k)
        Parameters:
            X : float or np.ndarray
                Real part
            Y : float or np.ndarray
                Imaginary part
        Returns:
            zeta : complex or np.ndarray
                Elliptic sine function result
        """
        if X is None or Y is None:
            raise ValueError(f"Invalid input to _SCT_g: X={X}, Y={Y} (cannot be None)")
        use_X = X / self.C0
        use_Y = Y / self.C0
        sn_X_k0, cn_X_k0, dn_X_k0, _ = sp.ellipj(use_X, self.m0)
        sn_Y_k0, cn_Y_k0, dn_Y_k0, _ = sp.ellipj(use_Y, self.m0)
        sn_X_k0p, cn_X_k0p, dn_X_k0p, _ = sp.ellipj(use_X, self.m0p)
        sn_Y_k0p, cn_Y_k0p, dn_Y_k0p, _ = sp.ellipj(use_Y, self.m0p)
        num = sn_X_k0 * dn_Y_k0p + 1j * sn_Y_k0p * cn_Y_k0p * cn_X_k0 * dn_X_k0
        den = cn_Y_k0p ** 2 + self.k0 ** 2 * sn_X_k0 ** 2 * sn_Y_k0p ** 2
        return num / den

    def _SCT_g_prime(self, X, Y):
        """
        计算椭圆正弦函数 sn(z, k) 的导数 cn(z, k) * dn(z, k)
        Parameters:
            X : float or np.ndarray
                Real part
            Y : float or np.ndarray
                Imaginary part
        Returns:
            dsndZ : complex or np.ndarray
                Derivative of elliptic sine function
        """
        if X is None or Y is None:
            raise ValueError(f"Invalid input to _SCT_g_prime: X={X}, Y={Y} (cannot be None)")
        use_X = X / self.C0
        use_Y = Y / self.C0
        sn_X_k0, cn_X_k0, dn_X_k0, _ = sp.ellipj(use_X, self.m0)
        sn_Y_k0, cn_Y_k0, dn_Y_k0, _ = sp.ellipj(use_Y, self.m0)
        sn_X_k0p, cn_X_k0p, dn_X_k0p, _ = sp.ellipj(use_X, self.m0p)
        sn_Y_k0p, cn_Y_k0p, dn_Y_k0p, _ = sp.ellipj(use_Y, self.m0p)
        sn_iY_k0 = 1j * (sn_Y_k0p / cn_Y_k0p)
        cn_iY_k0 = 1 / cn_Y_k0p
        dn_iY_k0 = dn_Y_k0p / cn_Y_k0p
        CN_num = sn_X_k0 * cn_X_k0 * dn_iY_k0 - sn_iY_k0 * cn_iY_k0 * dn_X_k0
        CN_den = sn_X_k0 * cn_iY_k0 * dn_iY_k0 - sn_iY_k0 * cn_X_k0 * dn_X_k0
        CN_Z_k0 = CN_num / CN_den
        DN_num = dn_X_k0 * cn_Y_k0p * dn_Y_k0p - 1j * self.m0 * sn_X_k0 * cn_X_k0 * sn_Y_k0p
        DN_den = cn_Y_k0p ** 2 + self.m0 * sn_X_k0 ** 2 * sn_Y_k0 ** 2
        DN_Z_k0 = DN_num / DN_den
        return CN_Z_k0 * DN_Z_k0 / self.C0  # 考虑归一化因子

    def cplx_W(self, X, Y):
        """
        计算复位势及其分量
        Parameters:
            X : np.ndarray
                Cartesian X coordinates
            Y : np.ndarray
                Cartesian Y coordinates
        Returns:
            w_ptn : np.ndarray
                Complex potential
            phi : np.ndarray
                Potential function
            psi : np.ndarray
                Stream function
        """
        if self.z_so is None or self.z_si is None:
            self.compute_z_coordinates()
        if self.z_so is None or self.z_si is None:
            raise ValueError("z_so or z_si is None after compute_z_coordinates")

        Z = X + 1j * Y
        w_ptn = np.zeros_like(Z, dtype=np.complex128)
        phi = np.zeros_like(Z, dtype=np.float64)
        psi = np.zeros_like(Z, dtype=np.float64)

        zeta = self._SCT_g(X, Y)
        zeta1 = self._SCT_g(np.real(self.z_so), np.imag(self.z_so))
        zeta2 = self._SCT_g(np.real(self.z_si), np.imag(self.z_si))

        w_ptn = (self.Q_source / (2 * np.pi)) * np.log(zeta - zeta1 + self.epsilon) + \
                (self.Q_sink / (2 * np.pi)) * np.log(zeta - zeta2 + self.epsilon)
        phi = w_ptn.real
        psi = w_ptn.imag

        return w_ptn, phi, psi

    def cplx_velocity(self, X, Y):
        """
        计算复势 w 对复坐标 Z 的导数
        Parameters:
            X : np.ndarray
                Cartesian X coordinates
            Y : np.ndarray
                Cartesian Y coordinates
        Returns:
            magV : np.ndarray
                Velocity magnitude
            u_vel : np.ndarray
                Horizontal velocity component
            v_vel : np.ndarray
                Vertical velocity component
        """
        Z = X + 1j * Y
        magV = np.zeros_like(Z, dtype=np.float64)
        u_vel = np.zeros_like(Z, dtype=np.float64)
        v_vel = np.zeros_like(Z, dtype=np.float64)

        zeta = self._SCT_g(X, Y)
        zeta1 = self._SCT_g(np.real(self.z_so), np.imag(self.z_so))
        zeta2 = self._SCT_g(np.real(self.z_si), np.imag(self.z_si))
        dsndz = self._SCT_g_prime(X, Y)

        dwdZ = dsndz * ((self.Q_source / (2 * np.pi)) / (zeta - zeta1 + self.epsilon) +
                        (self.Q_sink / (2 * np.pi)) / (zeta - zeta2 + self.epsilon))

        magV = np.abs(dwdZ)
        u_vel = dwdZ.real
        v_vel = -dwdZ.imag

        return magV, u_vel, v_vel

    def pressure(self, U_ref, magV, u_vel, v_vel):
        """
        计算压力分布和压力系数 CP
        Parameters:
            U_ref : float
                Reference velocity
            magV : np.ndarray
                Velocity magnitude
            u_vel : np.ndarray
                Horizontal velocity component
            v_vel : np.ndarray
                Vertical velocity component
        Returns:
            p : np.ndarray
                Pressure distribution
            CP : np.ndarray
                Pressure coefficient distribution
            d_v2 : np.ndarray
                Velocity error term (for diagnostics)
        """
        if not (magV.shape == u_vel.shape == v_vel.shape):
            raise ValueError("Input arrays magV, u_vel, and v_vel must have the same shape")

        rho = 1.2  # 空气密度 (kg/m^3)
        V_V = u_vel ** 2 + v_vel ** 2
        d_v2 = magV ** 2 - V_V
        p = 0.5 * rho * (U_ref ** 2 - V_V)
        CP = p / (0.5 * rho * U_ref ** 2)- 0.99999  # 调整系数，需说明物理意义

        return p, CP, d_v2


"""____示例____"""
# Example of using the general models
if __name__ == "__main__":
    Yheight = 2.0  # m
    Xlength = 4.0  # m
    N = 50
    solver = ParameterSolver(length=Xlength, height=Yheight)
    print(f"输入：宽{solver.L}高{solver.H}")
    result = solver.solve()
    if result:
        C, k, k_prime = result
        print(f"最终结果：C={C}, k={k}, k'={k_prime}")
    grid_dimension = Grid(Xlength, Yheight, N, origin=(2.0, 0.0))  # 网格坐标(有量纲)
    plotter = GridPlotter(grid_dimension)
    #plotter.plot_grid()

